/***************************************************************************
 * Copyright (c) Wolf Vollprecht, Johan Mabille and Sylvain Corlay          *
 * Copyright (c) QuantStack                                                 *
 *                                                                          *
 * Distributed under the terms of the BSD 3-Clause License.                 *
 *                                                                          *
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/

// This file is generated from test/files/cppy_source/test_dot_extended.cppy by
// preprocess.py!

#include <algorithm>

#include "xtensor/xarray.hpp"
#include "xtensor/xtensor.hpp"

#include "gtest/gtest.h"
#include "xtensor-blas/xlinalg.hpp"

namespace xt
{
    using namespace xt::placeholders;

    /*py
    a = np.random.random((2, 3, 5))
    b = np.random.random((4, 5))
    dr = np.dot(a, b.T)
    */
    TEST(xtest_extended, dot_broadcast)
    {
        // py_a
        xarray<double> py_a = {
            {{0.3745401188473625, 0.9507143064099162, 0.7319939418114051, 0.5986584841970366, 0.1560186404424365},
             {0.1559945203362026, 0.0580836121681995, 0.8661761457749352, 0.6011150117432088, 0.7080725777960455},
             {0.0205844942958024, 0.9699098521619943, 0.8324426408004217, 0.2123391106782762, 0.1818249672071006}},

            {{0.1834045098534338, 0.3042422429595377, 0.5247564316322378, 0.4319450186421158, 0.2912291401980419},
             {0.6118528947223795, 0.1394938606520418, 0.2921446485352182, 0.3663618432936917, 0.4560699842170359},
             {0.7851759613930136, 0.1996737821583597, 0.5142344384136116, 0.5924145688620425, 0.0464504127199977}}};
        // py_b
        xarray<double> py_b = {
            {0.6075448519014384, 0.1705241236872915, 0.0650515929852795, 0.9488855372533332, 0.9656320330745594},
            {0.8083973481164611, 0.3046137691733707, 0.0976721140063839, 0.6842330265121569, 0.4401524937396013},
            {0.1220382348447788, 0.4951769101112702, 0.0343885211152184, 0.9093204020787821, 0.2587799816000169},
            {0.662522284353982, 0.311711076089411, 0.5200680211778108, 0.5467102793432796, 0.184854455525527}};
        // py_dr
        xarray<double> py_dr = {
            {{1.1560019913607258, 1.1421672030085086, 1.1263990512143978, 1.2813094834150083},
             {1.415151366639716, 0.9513625344824885, 0.807426629014782, 1.0314517921651605},
             {0.6091122748507029, 0.6187149240291543, 0.7515524775267591, 0.898595256683809}},

            {{0.8885299172713558, 0.7159304454839006, 0.6592223836380569, 0.7792380767202456},
             {1.2025508600129964, 1.0170636073271262, 0.6049520893427571, 0.8853834024749684},
             {1.15151820221699, 1.1715787914743192, 0.763094187597877, 1.182339688054495}}};

        xt::xtensor<double, 2> bas = xt::transpose(py_b);

        auto xres = xt::linalg::dot(py_a, xt::transpose(py_b));
        auto xres2 = xt::linalg::dot(py_a, bas);
        std::cout << xres << std::endl;
        EXPECT_TRUE(xt::allclose(xres, py_dr));
        EXPECT_TRUE(xt::allclose(xres2, py_dr));
    }

    /*py
    a = np.random.random((2, 3, 5))
    b = np.random.random((5))
    dr = np.dot(a, b)
    */
    TEST(xtest_extended, dot_broadcast_2)
    {
        // py_a
        xarray<double> py_a = {
            {{0.9695846277645586, 0.7751328233611146, 0.9394989415641891, 0.8948273504276488, 0.5978999788110851},
             {0.9218742350231168, 0.0884925020519195, 0.1959828624191452, 0.0452272889105381, 0.3253303307632643},
             {0.388677289689482, 0.2713490317738959, 0.8287375091519293, 0.3567533266935893, 0.2809345096873808}},

            {{0.5426960831582485, 0.1409242249747626, 0.8021969807540397, 0.0745506436797708, 0.9868869366005173},
             {0.7722447692966574, 0.1987156815341724, 0.0055221171236024, 0.8154614284548342, 0.7068573438476171},
             {0.7290071680409873, 0.7712703466859457, 0.0740446517340904, 0.3584657285442726, 0.1158690595251297}}};
        // py_b
        xarray<double> py_b = {
            0.8631034258755935,
            0.6232981268275579,
            0.3308980248526492,
            0.0635583502860236,
            0.3109823217156622};
        // py_dr
        xarray<double> py_dr = {
            {1.8736790686065976, 1.0197269167779506, 0.8888679673881792},
            {1.1333287572487494, 1.0638629967411402, 1.1932578950872312}};

        auto xres = xt::linalg::dot(py_a, py_b);
        std::cout << xres << std::endl;
        EXPECT_TRUE(xt::allclose(xres, py_dr));
    }

    /*py
    a = np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5)
    b = np.arange(4 * 5 * 3).reshape(4, 5, 3)
    dr = np.dot(a, b)
    */
    TEST(xtest_extended, dot_broadcast_3)
    {
        // py_a
        xarray<long> py_a = {
            {{{0, 1, 2, 3, 4}, {5, 6, 7, 8, 9}, {10, 11, 12, 13, 14}, {15, 16, 17, 18, 19}},

             {{20, 21, 22, 23, 24}, {25, 26, 27, 28, 29}, {30, 31, 32, 33, 34}, {35, 36, 37, 38, 39}},

             {{40, 41, 42, 43, 44}, {45, 46, 47, 48, 49}, {50, 51, 52, 53, 54}, {55, 56, 57, 58, 59}}},

            {{{60, 61, 62, 63, 64}, {65, 66, 67, 68, 69}, {70, 71, 72, 73, 74}, {75, 76, 77, 78, 79}},

             {{80, 81, 82, 83, 84}, {85, 86, 87, 88, 89}, {90, 91, 92, 93, 94}, {95, 96, 97, 98, 99}},

             {{100, 101, 102, 103, 104},
              {105, 106, 107, 108, 109},
              {110, 111, 112, 113, 114},
              {115, 116, 117, 118, 119}}}};
        // py_b
        xarray<long> py_b = {
            {{0, 1, 2}, {3, 4, 5}, {6, 7, 8}, {9, 10, 11}, {12, 13, 14}},

            {{15, 16, 17}, {18, 19, 20}, {21, 22, 23}, {24, 25, 26}, {27, 28, 29}},

            {{30, 31, 32}, {33, 34, 35}, {36, 37, 38}, {39, 40, 41}, {42, 43, 44}},

            {{45, 46, 47}, {48, 49, 50}, {51, 52, 53}, {54, 55, 56}, {57, 58, 59}}};
        // py_dr
        xarray<long> py_dr = {
            {{{{90, 100, 110}, {240, 250, 260}, {390, 400, 410}, {540, 550, 560}},

              {{240, 275, 310}, {765, 800, 835}, {1290, 1325, 1360}, {1815, 1850, 1885}},

              {{390, 450, 510}, {1290, 1350, 1410}, {2190, 2250, 2310}, {3090, 3150, 3210}},

              {{540, 625, 710}, {1815, 1900, 1985}, {3090, 3175, 3260}, {4365, 4450, 4535}}},

             {{{690, 800, 910}, {2340, 2450, 2560}, {3990, 4100, 4210}, {5640, 5750, 5860}},

              {{840, 975, 1110}, {2865, 3000, 3135}, {4890, 5025, 5160}, {6915, 7050, 7185}},

              {{990, 1150, 1310}, {3390, 3550, 3710}, {5790, 5950, 6110}, {8190, 8350, 8510}},

              {{1140, 1325, 1510}, {3915, 4100, 4285}, {6690, 6875, 7060}, {9465, 9650, 9835}}},

             {{{1290, 1500, 1710}, {4440, 4650, 4860}, {7590, 7800, 8010}, {10740, 10950, 11160}},

              {{1440, 1675, 1910}, {4965, 5200, 5435}, {8490, 8725, 8960}, {12015, 12250, 12485}},

              {{1590, 1850, 2110}, {5490, 5750, 6010}, {9390, 9650, 9910}, {13290, 13550, 13810}},

              {{1740, 2025, 2310}, {6015, 6300, 6585}, {10290, 10575, 10860}, {14565, 14850, 15135}}}},

            {{{{1890, 2200, 2510}, {6540, 6850, 7160}, {11190, 11500, 11810}, {15840, 16150, 16460}},

              {{2040, 2375, 2710}, {7065, 7400, 7735}, {12090, 12425, 12760}, {17115, 17450, 17785}},

              {{2190, 2550, 2910}, {7590, 7950, 8310}, {12990, 13350, 13710}, {18390, 18750, 19110}},

              {{2340, 2725, 3110}, {8115, 8500, 8885}, {13890, 14275, 14660}, {19665, 20050, 20435}}},

             {{{2490, 2900, 3310}, {8640, 9050, 9460}, {14790, 15200, 15610}, {20940, 21350, 21760}},

              {{2640, 3075, 3510}, {9165, 9600, 10035}, {15690, 16125, 16560}, {22215, 22650, 23085}},

              {{2790, 3250, 3710}, {9690, 10150, 10610}, {16590, 17050, 17510}, {23490, 23950, 24410}},

              {{2940, 3425, 3910}, {10215, 10700, 11185}, {17490, 17975, 18460}, {24765, 25250, 25735}}},

             {{{3090, 3600, 4110}, {10740, 11250, 11760}, {18390, 18900, 19410}, {26040, 26550, 27060}},

              {{3240, 3775, 4310}, {11265, 11800, 12335}, {19290, 19825, 20360}, {27315, 27850, 28385}},

              {{3390, 3950, 4510}, {11790, 12350, 12910}, {20190, 20750, 21310}, {28590, 29150, 29710}},

              {{3540, 4125, 4710}, {12315, 12900, 13485}, {21090, 21675, 22260}, {29865, 30450, 31035}}}}};

        auto xres = xt::linalg::dot(py_a, py_b);
        EXPECT_TRUE(xt::allclose(xres, py_dr));
    }
}  // namespace xt
